import sys
import rasterio
from rasterio.plot import show
from PyQt5.QtWidgets import QApplication, QWidget, QPushButton, QFileDialog, QLabel, QVBoxLayout, QComboBox, QLineEdit
from PyQt5.QtGui import QPixmap
from io import BytesIO
from PIL import Image
import numpy as np

def show_file_dialog():
    app = QApplication(sys.argv)
    window = QWidget()

    # 整体布局
    layout = QVBoxLayout()

    # 用于显示文件路径的文本框
    file_path_edit = QLineEdit(window)
    layout.addWidget(file_path_edit)

    # 用于显示遥感影像的标签
    image_label = QLabel(window)
    layout.addWidget(image_label)

    # 红色波段选择下拉框
    red_band_combobox = QComboBox(window)
    layout.addWidget(red_band_combobox)
    # 绿色波段选择下拉框
    green_band_combobox = QComboBox(window)
    layout.addWidget(green_band_combobox)
    # 蓝色波段选择下拉框
    blue_band_combobox = QComboBox(window)
    layout.addWidget(blue_band_combobox)

    # 触发文件选择操作的按钮
    choose_file_button = QPushButton('Choose Image', window)

    def handle_button_click():
        options = QFileDialog.Options()
        file_name, _ = QFileDialog.getOpenFileName(window, "Select Image", ".", "GeoTIFF Files (*.tif *.tiff);;Images (*.png *.jpg *.jpeg);;All Files (*)", options=options)
        if file_name:
            file_path_edit.setText(file_name)
            try:
                # 使用rasterio打开遥感影像
                with rasterio.open(file_name) as src:
                    # 获取波段数量
                    num_bands = src.count
                    # 清空波段选择下拉框并添加波段选项
                    red_band_combobox.clear()
                    green_band_combobox.clear()
                    blue_band_combobox.clear()
                    for i in range(1, num_bands + 1):
                        red_band_combobox.addItem(f'Band {i}')
                        green_band_combobox.addItem(f'Band {i}')
                        blue_band_combobox.addItem(f'Band {i}')
                    # 初始设置为真彩色（可根据需求调整）
                    red_band_combobox.setCurrentIndex(0)
                    green_band_combobox.setCurrentIndex(1)
                    blue_band_combobox.setCurrentIndex(2)
                    display_band_data(file_name)
            except Exception as e:
                print(f"Error opening or processing the file: {e}")

    def display_band_data(file_name):
        # 获取当前选择的波段索引
        red_band_index = red_band_combobox.currentIndex() + 1
        green_band_index = green_band_combobox.currentIndex() + 1
        blue_band_index = blue_band_combobox.currentIndex() + 1
        try:
            with rasterio.open(file_name) as src:
                # 读取指定波段的数据
                r_band = src.read(red_band_index)
                g_band = src.read(green_band_index)
                b_band = src.read(blue_band_index)
                # 将各波段数据转换为合适的格式用于显示（这里简单处理为8位无符号整数）
                # 注意：这只是一个简单的转换示例，实际可能需要更复杂的处理
                r_band = ((r_band - r_band.min()) * (255.0 / (r_band.max() - r_band.min()))).astype('uint8')
                g_band = ((g_band - g_band.min()) * (255.0 / (g_band.max() - g_band.min()))).astype('uint8')
                b_band = ((b_band - b_band.min()) * (255.0 / (b_band.max() - b_band.min()))).astype('uint8')
                # 合并RGB三个波段数据为一个三维数组，用于构建彩色图像
                img_data = np.dstack((r_band, g_band, b_band))
                # 使用PIL将数组转换为图像对象
                img = Image.fromarray(img_data, 'RGB')
                # 将图像对象保存到内存文件对象中
                output = BytesIO()
                img.save(output, 'PNG')
                output.seek(0)
                # 使用QPixmap加载内存中的图像数据并设置到标签上
                pixmap = QPixmap()
                pixmap.loadFromData(output.read())
                image_label.setPixmap(pixmap.scaled(image_label.size(), aspectRatioMode=1))
        except Exception as e:
            print(f"Error opening or processing the file: {e}")

    def handle_red_band_change(index):
        display_band_data(file_path_edit.text())

    def handle_green_band_change(index):
        display_band_data(file_path_edit.text())

    def handle_blue_band_change(index):
        display_band_data(file_path_edit.text())

    choose_file_button.clicked.connect(handle_button_click)
    # 当红色波段选择下拉框的选项改变时触发
    red_band_combobox.currentIndexChanged.connect(handle_red_band_change)
    # 当绿色波段选择下拉框的选项改变时触发
    green_band_combobox.currentIndexChanged.connect(handle_green_band_change)
    # 当蓝色波段选择下拉框的选项改变时触发
    blue_band_combobox.currentIndexChanged.connect(handle_blue_band_change)

    layout.addWidget(choose_file_button)

    window.setLayout(layout)
    window.setGeometry(200, 50, 600, 600)  # 适当增加窗口高度以容纳更多下拉框
    window.setWindowTitle('Remote Sensing Image Viewer')
    window.show()
    sys.exit(app.exec_())


if __name__ == '__main__':
    show_file_dialog()